# -*- coding: utf-8 -*-
import pandas as pd
import talib as ta
import numpy as np
import matplotlib.pyplot as plt
from py_mysql import *
from tqdm import tqdm

csv_arr = []
def run_all(code,startDate,endDate):
    obj_data = {
        'code': code,
        'score': None,
        'start_date': startDate,
        'end_date': endDate,
        'y_pred_proba': None,
        'importances_sort_df': None,
        'learning_rate':None,
        'max_depth':None,
        'n_estimators':None
    }
    startDate = startDate  #开始时间
    endDate = endDate   #结束时间
    codeArr = [code]   #品种数组
    query_db = Mysql_search()
    df = query_db.get_one(codeArr,startDate,endDate)
    df = df[codeArr[0]]

    timep = 10
    timp_RSI = 5
    df['RSI'] = ta.𝑅𝑆𝐼(df['close'], 𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑=timp_RSI)

    # 一ADX
    df['MINUSDI'] = ta.𝑀𝐼𝑁𝑈𝑆_𝐷𝐼(df['high'], df['low'], df['close'], 𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑=timep)
    df['PLUSDI'] = ta.𝑃𝐿𝑈𝑆_𝐷𝐼(df['high'], df['low'], df['close'], 𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑=timep)
    df['ADXR'] = ta.ADXR(df['high'], df['low'], df['close'], timeperiod=timep)
    df['ADX_shift'] = df['PLUSDI'].shift(1) - df['MINUSDI'].shift(1)
    def ADX_status_edit(val,y):
        if val['ADX_shift'] < 0 and val['PLUSDI'] > val['MINUSDI'] and val['ADXR'] > y:
            return 1
        elif val['ADX_shift'] > 0 and val['PLUSDI'] < val['MINUSDI'] and val['ADXR'] > y:
            return -1
        else:
            return -1
    df['ADX_25'] = df.apply(lambda x: ADX_status_edit(x,25), axis=1)
    df['ADX_55'] = df.apply(lambda x: ADX_status_edit(x,50), axis=1)
    df['ADX_75'] = df.apply(lambda x: ADX_status_edit(x,75), axis=1)
    df = df.drop(columns=['MINUSDI','PLUSDI','ADXR','ADX_shift'])

    # 二APO
    df['APO'] = ta.APO(df['close'], fastperiod=5, slowperiod=10, matype=0)
    df['APO_ch'] = df['APO'].apply(lambda x: 1 if x > 0 else ( -1 if x < 0 else 0))
    df = df.drop(columns=['APO'])
    
    # 三AROON
    UAROON,DAROON =ta.AROON(df['high'], df['low'], timeperiod=timep)
    df['UAROON'] = UAROON
    df['DAROON'] = DAROON
    df['UAROON_bp'] = df['UAROON'] - df['UAROON'].shift(1)
    df['DAROON_bp'] = df['DAROON'] - df['DAROON'].shift(1)
    def AROON_edit(val):
        if (val["UAROON_bp"] > 0 and abs(val["UAROON"] - 100) < 5) or (val['DAROON_bp'] < 0 and abs(val['DAROON'] - 0) < 5):
            return 1
        elif (val["UAROON_bp"] < 0 and abs(val["UAROON"] - 100) < 5) or (val['DAROON_bp'] >= 0 and abs(val['DAROON'] - 0) < 5):
            return -1
        else:
            return 0
    df['AROON_ch'] = df.apply(lambda x: AROON_edit(x), axis=1)

    df['AROON_shift'] = df['UAROON'].shift(1) - df['DAROON'].shift(1)
    def AROON_status_edit(val):
        if val['AROON_shift'] < 0 and val['UAROON'] > val['DAROON']:
            return 1
        elif val['AROON_shift'] > 0 and val['UAROON'] < val['DAROON']:
            return -1
        else:
            return 0
    df['AROON_up_down'] = df.apply(lambda x: AROON_status_edit(x), axis=1)
    df = df.drop(columns=['UAROON','DAROON','UAROON_bp','DAROON_bp','AROON_shift'])

    # 四BOP
    df['BOP'] = ta.BOP(df['open'],df['high'],df['low'],df['close'])
    df['BOP_ch'] = df['BOP'].apply(lambda x: 1 if x > 0 else (-1 if x < 0 else 0))
    def BOP_edit(val):
        if val['BOP'] > 0 and val['RSI'] < 20:
            return 1
        elif val['BOP'] < 0 and val['RSI'] > 80:
            return -1
        else:
            return 0
    df['BOP_and_RSI'] = df.apply(lambda x: BOP_edit(x), axis=1)
    df = df.drop(columns=['BOP'])
    
    # 五CCI
    df['CCI'] = ta.CCI(df['high'],df['low'],df['close'],𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑=timep)
    df['CCI_ch'] = df['CCI'].apply(lambda x: -1 if x > 100 else (1 if x < -100 else 0))
    df['willr'] = ta.WILLR(df['high'],df['low'],df['close'], 𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑=timep)
    def CCI_edit(val):
        if val['CCI'] < -100 and val['willr'] > 80:
            return 1
        elif val['CCI'] > 100 and val['willr'] < 20:
            return -1
        else:
            return 0
    df['CCI_and_willr'] = df.apply(lambda x: CCI_edit(x), axis=1)
    df = df.drop(columns=['CCI','willr'])

    # 六CMO
    df['CMO_time_5'] = ta.CMO(df['close'],𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑=5)
    df['CMO_5'] = df['CMO_time_5'].apply(lambda x: -1 if x > 50 else (1 if x < -50 else 0))
    df['CMO_time_10'] = ta.CMO(df['close'],𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑=10)
    df['CMO_10'] = df['CMO_time_10'].apply(lambda x: -1 if x > 50 else (1 if x < -50 else 0))
    df['CMO_time_15'] = ta.CMO(df['close'],𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑=15)
    df['CMO_15'] = df['CMO_time_15'].apply(lambda x: -1 if x > 50 else (1 if x < -50 else 0))
    df = df.drop(columns=['CMO_time_5','CMO_time_10','CMO_time_15'])

    #七MACD
    macd, macdsignal, macdhist = ta.MACD(df['close'], 𝑓𝑎𝑠𝑡𝑝𝑒𝑟𝑖𝑜𝑑=5, 𝑠𝑙𝑜𝑤𝑝𝑒𝑟𝑖𝑜𝑑=10, 𝑠𝑖𝑔𝑛𝑎𝑙𝑝𝑒𝑟𝑖𝑜𝑑=7)
    df['macd'] = macd
    df['macdsignal'] = macdsignal
    df['macd_macdsignal_up_down'] = df['macd'].shift(1) - df['macdsignal'].shift(1)
    def macd_edit(val):
        if val['macd'] > 0 and val['macdsignal'] > 0 and val['macd_macdsignal_up_down'] < 0 and val['macd'] > val['macdsignal']:
            return 1
        elif val['macd'] < 0 and val['macdsignal'] < 0 and val['macd_macdsignal_up_down'] > 0 and val['macd'] < val['macdsignal']:
            return -1
        else:
            return 0
    df['macd_ch'] = df.apply(lambda x:macd_edit(x), axis=1)
    def macd_RSI_edit(val):
        if val['macd'] > 0 and val['macdsignal'] > 0 and val['macd_macdsignal_up_down'] < 0 and val['macd'] > val['macdsignal'] and val['RSI'] < 20:
            return 1
        elif val['macd'] < 0 and val['macdsignal'] < 0 and val['macd_macdsignal_up_down'] > 0 and val['macd'] < val['macdsignal'] and val['RSI'] > 80:
            return -1
        else:
            return 0
    df['macd_RSI_ch'] = df.apply(lambda x:macd_RSI_edit(x), axis=1)
    df = df.drop(columns=['macd','macdsignal','macd_macdsignal_up_down'])
    #MACDFIX
    macdfix, macdsignalfix, macdhistfix = ta.MACDFIX(df['close'], 𝑠𝑖𝑔𝑛𝑎𝑙𝑝𝑒𝑟𝑖𝑜𝑑=9)
    df['macdfix'] = macdfix
    df['macdsignalfix'] = macdsignalfix
    df['macdfix_macdsignalfix_up_down'] = df['macdfix'].shift(1) - df['macdsignalfix'].shift(1)
    def macdfix_edit(val):
        if val['macdfix'] > 0 and val['macdsignalfix'] > 0 and val['macdfix_macdsignalfix_up_down'] < 0 and val['macdfix'] > val['macdsignalfix']:
            return 1
        elif val['macdfix'] < 0 and val['macdsignalfix'] < 0 and val['macdfix_macdsignalfix_up_down'] > 0 and val['macdfix'] < val['macdsignalfix']:
            return -1
        else:
            return 0
    df['macdfix_ch'] = df.apply(lambda x:macdfix_edit(x), axis=1)
    def macdfix_RSI_edit(val):
        if val['macdfix'] > 0 and val['macdsignalfix'] > 0 and val['macdfix_macdsignalfix_up_down'] < 0 and val['macdfix'] > val['macdsignalfix'] and val['RSI'] < 20:
            return 1
        elif val['macdfix'] < 0 and val['macdsignalfix'] < 0 and val['macdfix_macdsignalfix_up_down'] > 0 and val['macdfix'] < val['macdsignalfix'] and val['RSI'] > 80:
            return -1
        else:
            return 0
    df['macdfix_RSI_ch'] = df.apply(lambda x:macdfix_RSI_edit(x), axis=1)
    df = df.drop(columns=['macdfix','macdsignalfix','macdfix_macdsignalfix_up_down'])

    # 八MFI
    df['MFI'] = ta.𝑀𝐹𝐼(df['high'],df['low'],df['close'],df['volume'],𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑=timep)
    df['MFI_shift'] = df['MFI'].shift(1)
    def MFI_edit(val):
        if ((val['MFI_shift'] > 80 and val['MFI'] < 80) or (val['MFI_shift'] < 20 and val['MFI'] > 20)) and val['RSI'] > 80:
            return -1
        elif ((val['MFI_shift'] > 80 and val['MFI'] < 80) or (val['MFI_shift'] < 20 and val['MFI'] > 20)) and val['RSI'] < 20:
            return 1
        else:
            return 0
    df['MFI_and_RSI'] = df.apply(lambda x:MFI_edit(x), axis=1)
    df = df.drop(columns=['MFI','MFI_shift'])

    # 九 MOM
    df['MOM_time_5'] = ta.𝑀𝑂𝑀(df['close'], 𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑=5)
    df['MOM_5'] = df['MOM_time_5'].apply(lambda x: -1 if x > 1 else (1 if x < -1 else 0))
    df['MOM_time_10'] = ta.𝑀𝑂𝑀(df['close'], 𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑=10)
    df['MOM_10'] = df['MOM_time_10'].apply(lambda x: -1 if x > 1 else (1 if x < -1 else 0))
    df['MOM_time_25'] = ta.𝑀𝑂𝑀(df['close'], 𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑=25)
    df['MOM_25'] = df['MOM_time_25'].apply(lambda x: 1 if x > 0 else (-1 if x < 0 else 0))
    df = df.drop(columns=['MOM_time_5','MOM_time_10','MOM_time_25'])

    # 十 PPO
    df['PPO_time_5'] = ta.𝑃𝑃𝑂(df['close'], 𝑓𝑎𝑠𝑡𝑝𝑒𝑟𝑖𝑜𝑑=5, 𝑠𝑙𝑜𝑤𝑝𝑒𝑟𝑖𝑜𝑑=10, matype=0)
    df['PPO_5'] = df['PPO_time_5'].apply(lambda x: 1 if x > 0 else (-1 if x < 0 else 0))
    df['PPO_time_10'] = ta.𝑃𝑃𝑂(df['close'], 𝑓𝑎𝑠𝑡𝑝𝑒𝑟𝑖𝑜𝑑=10, 𝑠𝑙𝑜𝑤𝑝𝑒𝑟𝑖𝑜𝑑=20, matype=0)
    df['PPO_10'] = df['PPO_time_10'].apply(lambda x: 1 if x > 0 else (-1 if x < 0 else 0))
    df = df.drop(columns=['PPO_time_5','PPO_time_10'])

    # 十一 RSI
    df['RSI_short'] = ta.𝑅𝑆𝐼(df['close'], 𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑=5)
    df['RSI_long'] = ta.𝑅𝑆𝐼(df['close'], 𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑=10)
    df['RSI_short_up_down'] = df['RSI_short'].shift(1) - df['RSI_long'].shift(1)
    df['RSI_short_shift'] = df['RSI_short'].shift(1)
    def RSI_edit(val):
        if val['RSI_short'] > val['RSI_short_shift'] and val['RSI_short'] > val['RSI_long'] and val['RSI_short_up_down'] < 0:
            return 1 
        elif val['RSI_short'] < val['RSI_short_shift'] and val['RSI_short'] < val['RSI_long'] and val['RSI_short_up_down'] > 0:
            return -1
        else:
            return 0
    df['RSI_ch'] = df.apply(lambda x: RSI_edit(x), axis=1)

    df = df.drop(columns=['RSI_short','RSI_long','RSI_short_up_down','RSI_short_shift'])

    # 十二 KDJ
    slowk, slowd = ta.𝑆𝑇𝑂𝐶𝐻(df['high'],df['low'],df['close'], 𝑓𝑎𝑠𝑡𝑘_𝑝𝑒𝑟𝑖𝑜𝑑=5, 𝑠𝑙𝑜𝑤𝑘_𝑝𝑒𝑟𝑖𝑜𝑑=10, 𝑠𝑙𝑜𝑤𝑘_𝑚𝑎𝑡𝑦𝑝𝑒=0, 𝑠𝑙𝑜𝑤𝑑_𝑝𝑒𝑟𝑖𝑜𝑑=3, 𝑠𝑙𝑜𝑤𝑑_𝑚𝑎𝑡𝑦𝑝𝑒=0)
    df['slowk']=slowk
    df['slowd']=slowd
    df['KDJ_ch_slowk_1'] = df['slowk'].apply(lambda x: -1 if x > 90 else (1 if x < 10 else 0))
    df['KDJ_ch_slowd_1'] = df['slowd'].apply(lambda x: -1 if x > 80 else (1 if x < 20 else 0))
    def KDJ_edit(val):
        if val['slowk'] > 90 and val['slowd'] > 80:
            return -1
        elif val['slowk'] < 10 and val['slowd'] < 20:
            return 1
        else:
            return 0
    df['KDJ_ch_2'] = df.apply(lambda x: KDJ_edit(x), axis=1)
    # 𝑆𝑇𝑂𝐶𝐻𝑅𝑆𝐼
    fastkRSI, fastdRSI = ta.𝑆𝑇𝑂𝐶𝐻𝑅𝑆𝐼(df['close'], 𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑=timep, 𝑓𝑎𝑠𝑡𝑘_𝑝𝑒𝑟𝑖𝑜𝑑=5, 𝑓𝑎𝑠𝑡𝑑_𝑝𝑒𝑟𝑖𝑜𝑑=3, 𝑓𝑎𝑠𝑡𝑑_𝑚𝑎𝑡𝑦𝑝𝑒=0)
    df['fastkRSI']= fastkRSI
    df['fastdRSI']= fastdRSI
    df['KDJ_ch_fastkRSI_1'] = df['fastkRSI'].apply(lambda x: -1 if x > 90 else (1 if x < 10 else 0))
    df['KDJ_ch_fastdRSI_1'] = df['fastdRSI'].apply(lambda x: -1 if x > 80 else (1 if x < 20 else 0))
    def STOCHKDJ_edit(val):
        if val['fastkRSI'] > 90 and val['fastdRSI'] > 80:
            return -1
        elif val['fastkRSI'] < 10 and val['fastdRSI'] < 20:
            return 1
        else:
            return 0
    df['STOCHKDJ_ch_2'] = df.apply(lambda x: STOCHKDJ_edit(x), axis=1)
    df = df.drop(columns=['slowk','slowd','fastkRSI','fastdRSI'])

    # 十三 UOS
    df['UOS'] = ta.𝑈𝐿𝑇𝑂𝑆𝐶(df['high'],df['low'],df['close'], 𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑1=5, 𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑2=10, 𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑3=15)
    df['UOS_up_short'] = df['UOS'].shift(1) - 50
    df['UOS_down_short'] = df['UOS'].shift(1) - 65
    def UOS_short_edit(val):
        if val['UOS_up_short'] < 0 and val['UOS'] > 50:
            return 1
        elif val['UOS_down_short'] > 0 and val['UOS'] < 65:
            return -1
        else:
            return 0
    df['UOS_short'] = df.apply(lambda x:UOS_short_edit(x), axis=1)
    df['UOS_up_long'] = df['UOS'].shift(1) - 35
    df['UOS_down_long'] = df['UOS'].shift(1) - 70
    def UOS_long_edit(val):
        if val['UOS_up_long'] < 0 and val['UOS'] > 35:
            return 1
        elif val['UOS_down_long'] > 0 and val['UOS'] < 70:
            return -1
        else:
            return 0
    df['UOS_long'] = df.apply(lambda x:UOS_long_edit(x), axis=1)
    df = df.drop(columns=['UOS','UOS_up_short','UOS_down_short','UOS_up_long','UOS_down_long'])
    
    # 十四 WR
    df['WR_time_5'] = ta.𝑊𝐼𝐿𝐿𝑅(df['high'],df['low'],df['close'], 𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑=5)
    df['WR_5'] = df['WR_time_5'].apply(lambda x: 1 if x > 80 else (-1 if x < 20 else 0))
    df['WR_time_10'] = ta.𝑊𝐼𝐿𝐿𝑅(df['high'],df['low'],df['close'], 𝑡𝑖𝑚𝑒𝑝𝑒𝑟𝑖𝑜𝑑=10)
    df['WR_10'] = df['WR_time_10'].apply(lambda x: 1 if x > 80 else (-1 if x < 20 else 0))
    df = df.drop(columns=['WR_time_5','WR_time_10'])

    df['close_ch'] =df['close'].shift(-1) - df['close']
    def load_bp(val):
        if val > 0:
            return 1
        elif val < 0:
            return -1
        else:
            return 0
    df['close_ch'] = df['close_ch'].map(load_bp)

    print(df)

    # 提取特征变量和目标变量
    x = df.drop(columns=['close','open','high','low','volume','close_ch','RSI'])
    y = df['close_ch']

    # 划分训练集和测试集
    from sklearn.model_selection import train_test_split
    X_train, X_test, y_train, y_test = train_test_split(x.values,y, test_size=0.2, random_state=123)

    # 模型训练及搭建
    from xgboost import XGBClassifier
    # clf = XGBClassifier(n_estimators=100, learning_rate=0.05)
    clf = XGBClassifier(n_estimators=50, learning_rate=0.01)
    clf.fit(X_train, y_train)

    # 模型预测及评估**
    y_pred = clf.predict(X_test)
    a = pd.DataFrame()  # 创建一个空DataFrame
    a['预测值'] = list(y_pred)
    a['实际值'] = list(y_test)

    from sklearn.metrics import accuracy_score
    score = accuracy_score(y_pred, y_test)
    obj_data['score'] = score
    # print(score)

    y_pred_proba = clf.predict_proba(X_test)
    obj_data['y_pred_proba'] = y_pred_proba[0:5]
    # print(y_pred_proba[0:5])  # 查看前5个预测的概率

    features = x.columns  # 获取特征名称
    importances = clf.feature_importances_  # 获取特征重要性
    importances_df = pd.DataFrame()
    importances_df['特征名称'] = features
    importances_df['特征重要性'] = importances
    importances_sort_df = importances_df.sort_values('特征重要性', ascending=False)
    obj_data['importances_sort_df'] = importances_sort_df[:20]
    # print(importances_sort_df[:20])

    # 模型参数调优**
    from sklearn.model_selection import GridSearchCV
    parameters = {'max_depth': [1, 3, 5, 7, 10], 'n_estimators': [50, 100, 150, 200, 250], 'learning_rate': [0.01, 0.05, 0.1, 0.2, 0.4, 0.8]}  # 指定模型中参数的范围
    clf = XGBClassifier()  # 构建模型
    grid_search = GridSearchCV(clf, parameters, scoring='roc_auc', cv=5)  

    print(y_train)
    grid_search.fit(X_train, y_train)  # 传入数据

    optimal_num = grid_search.best_params_
    obj_data['learning_rate'] = optimal_num['learning_rate']
    obj_data['max_depth'] = optimal_num['max_depth']
    obj_data['n_estimators'] = optimal_num['n_estimators']
    print(grid_search.best_params_)  # 输出参数的最优值


    csv_arr.append(obj_data)
    
    

if __name__ == '__main__':
    startDate = '2021-06-01'  #开始时间
    endDate = '2021-09-15'   #结束时间
    codeArr = ['AG', 'A', 'AL', 'AP', 'AU',  'BC', 'B', 'BU', 'CF', 'CJ', 'C', 'CS', 'CU', 'CY', 'EB', 'EG', 'FB', 'FG', 'FU', 'HC', 'IH', 'I', 'JD', 'J', 'JM', 'LH', 'L', 'LU', 'MA', 'M', 'NI', 'NR', 'OI', 'PB', 'PF', 'PG', 'PK', 'P', 'PP', 'RB', 'RR', 'RU', 'SA', 'SC', 'SF', 'SM', 'SN', 'SP', 'SR', 'SS', 'TA', 'TF', 'T', 'TS', 'UR', 'V', 'Y', 'ZN']   #品种数组
    codeArr = ['TS']
    for item in tqdm(codeArr):
        test = run_all(item,startDate,endDate)
    # print(csv_arr)
    # df = pd.DataFrame(csv_arr)
    # df.to_csv('importances_sort_df.csv',sep=',',index=True,header=True)


